"""
统一的分割评价指标和损失函数
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class DiceLoss(nn.Module):
    """
    Dice Loss
    用于处理类别不平衡问题，特别适合稀疏目标分割
    """

    def __init__(self, smooth=1e-6):
        super(DiceLoss, self).__init__()
        self.smooth = smooth

    def forward(self, pred, target):
        """
        Args:
            pred: 预测结果 (B, 1, H, W)，logits（未经sigmoid）
            target: 真实标签 (B, 1, H, W)，值为0或1

        Returns:
            loss: Dice损失
        """
        # 应用sigmoid
        pred = torch.sigmoid(pred)

        # 展平
        pred = pred.view(-1)
        target = target.view(-1)

        # 计算交集
        intersection = (pred * target).sum()

        # 计算Dice系数
        dice = (2.0 * intersection + self.smooth) / (pred.sum() + target.sum() + self.smooth)

        # Dice Loss = 1 - Dice
        return 1 - dice


class SegmentationMetrics:
    """
    分割任务的综合评价指标
    包括：IoU, Dice, Pixel Accuracy, Precision, Recall
    """

    def __init__(self, threshold=0.5, smooth=1e-6):
        self.threshold = threshold
        self.smooth = smooth
        self.reset()

    def reset(self):
        """重置所有累计值"""
        self.tp = 0  # True Positive
        self.fp = 0  # False Positive
        self.tn = 0  # True Negative
        self.fn = 0  # False Negative

    def update(self, pred, target):
        """
        更新混淆矩阵

        Args:
            pred: 预测结果 (B, 1, H, W)，值在[0, 1]之间（经过sigmoid）
            target: 真实标签 (B, 1, H, W)，值为0或1
        """
        if isinstance(pred, torch.Tensor):
            pred = pred.detach().cpu().numpy()
        if isinstance(target, torch.Tensor):
            target = target.detach().cpu().numpy()

        # 展平
        pred = pred.flatten()
        target = target.flatten()

        # 二值化
        pred = (pred > self.threshold).astype(np.float32)

        # 更新混淆矩阵
        self.tp += np.sum((pred == 1) & (target == 1))
        self.fp += np.sum((pred == 1) & (target == 0))
        self.tn += np.sum((pred == 0) & (target == 0))
        self.fn += np.sum((pred == 0) & (target == 1))

    def get_iou(self):
        """计算IoU（前景类）"""
        intersection = self.tp
        union = self.tp + self.fp + self.fn
        iou = (intersection + self.smooth) / (union + self.smooth)
        return iou

    def get_miou(self):
        """计算mIoU（背景+前景的平均）"""
        # 前景IoU
        fg_iou = self.tp / (self.tp + self.fp + self.fn + self.smooth)
        # 背景IoU
        bg_iou = self.tn / (self.tn + self.fp + self.fn + self.smooth)
        miou = (fg_iou + bg_iou) / 2.0
        return miou

    def get_dice(self):
        """计算Dice系数"""
        dice = (2.0 * self.tp + self.smooth) / (2.0 * self.tp + self.fp + self.fn + self.smooth)
        return dice

    def get_pixel_accuracy(self):
        """计算像素准确率 (PA)"""
        pa = (self.tp + self.tn + self.smooth) / (self.tp + self.tn + self.fp + self.fn + self.smooth)
        return pa

    def get_precision(self):
        """计算精确率（前景类）"""
        precision = (self.tp + self.smooth) / (self.tp + self.fp + self.smooth)
        return precision

    def get_recall(self):
        """计算召回率（前景类）"""
        recall = (self.tp + self.smooth) / (self.tp + self.fn + self.smooth)
        return recall

    def get_f1_score(self):
        """计算F1分数（等同于Dice）"""
        precision = self.get_precision()
        recall = self.get_recall()
        f1 = (2.0 * precision * recall + self.smooth) / (precision + recall + self.smooth)
        return f1

    def get_all_metrics(self):
        """获取所有指标"""
        return {
            'iou': self.get_iou(),
            'miou': self.get_miou(),
            'dice': self.get_dice(),
            'pixel_accuracy': self.get_pixel_accuracy(),
            'precision': self.get_precision(),
            'recall': self.get_recall(),
            'f1_score': self.get_f1_score()
        }


class CombinedLoss(nn.Module):
    """
    组合损失：BCE + Dice Loss
    BCE处理整体分类，Dice处理类别不平衡
    """

    def __init__(self, bce_weight=1.0, dice_weight=1.0):
        """
        Args:
            bce_weight: BCE损失的权重
            dice_weight: Dice损失的权重
        """
        super(CombinedLoss, self).__init__()
        self.bce_weight = bce_weight
        self.dice_weight = dice_weight
        self.bce_loss = nn.BCEWithLogitsLoss()
        self.dice_loss = DiceLoss()

    def forward(self, pred, target):
        """
        Args:
            pred: 预测结果 (B, 1, H, W)，logits
            target: 真实标签 (B, 1, H, W)，值为0或1

        Returns:
            loss: 组合损失
        """
        bce = self.bce_loss(pred, target)
        dice = self.dice_loss(pred, target)

        # 组合损失
        # 原因：BCE提供像素级监督，Dice缓解前景稀疏问题
        total_loss = self.bce_weight * bce + self.dice_weight * dice

        return total_loss, bce, dice
